import torch.nn as nn
import torch

class DNNModel(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.user_embedding = nn.Embedding(config.NUM_USERS, config.EMBEDDING_DIM)
        self.item_embedding = nn.Embedding(config.NUM_ITEMS, config.EMBEDDING_DIM)
        
        self.dnn = nn.Sequential(
            nn.Linear(config.EMBEDDING_DIM * 2, 128),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, 1),
            nn.Sigmoid()
        )
    
    def forward(self, user_ids, item_ids):
        user_embeds = self.user_embedding(user_ids)
        item_embeds = self.item_embedding(item_ids)
        concat_embeds = torch.cat([user_embeds, item_embeds], dim=1)
        return self.dnn(concat_embeds).squeeze()